#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging
import json
from aliyunsdkcore.client import AcsClient
from aliyunsdkcore.acs_exception.exceptions import ClientException, ServerException
from aliyunsdkswas_open.request.v20200601.ListFirewallRulesRequest import ListFirewallRulesRequest
from aliyunsdkswas_open.request.v20200601.CreateFirewallRuleRequest import CreateFirewallRuleRequest
from aliyunsdkswas_open.request.v20200601.DeleteFirewallRuleRequest import DeleteFirewallRuleRequest

logger = logging.getLogger(__name__)

class AlibabaCloudSWASManager:
    def __init__(self, access_key_id: str, access_key_secret: str, region_id: str):
        """
        初始化阿里云轻量应用服务器管理器
        
        Args:
            access_key_id: 阿里云访问密钥ID
            access_key_secret: 阿里云访问密钥
            region_id: 区域ID
        """
        self.client = AcsClient(access_key_id, access_key_secret, region_id)
        self.region_id = region_id
        
    def get_existing_rules(self, instance_id: str, description: str) -> list:
        """
        获取指定轻量应用服务器中由本程序添加的防火墙规则
        
        Args:
            instance_id: 实例ID
            description: 规则描述标识
            
        Returns:
            匹配的规则列表
        """
        try:
            request = ListFirewallRulesRequest()
            request.set_InstanceId(instance_id)
            request.set_RegionId(self.region_id)
            request.set_PageSize(100)  # 获取尽可能多的规则
            request.set_accept_format('json')
            
            response = self.client.do_action_with_exception(request)
            result = json.loads(response)
            
            matched_rules = []
            if 'FirewallRules' in result:
                for rule in result['FirewallRules']:
                    # 阿里云轻量应用服务器没有描述字段，我们使用备注字段来标识
                    # 我们将在创建规则时把描述作为备注保存
                    if rule.get('Remark') == description:
                        matched_rules.append(rule)
                        
            return matched_rules
        except Exception as e:
            logger.error(f"获取轻量应用服务器防火墙规则失败: {e}")
            return []
            
    def remove_rules(self, instance_id: str, rules: list):
        """
        删除指定的防火墙规则
        
        Args:
            instance_id: 实例ID
            rules: 要删除的规则列表
        """
        for rule in rules:
            try:
                request = DeleteFirewallRuleRequest()
                request.set_InstanceId(instance_id)
                request.set_RegionId(self.region_id)
                request.set_RuleId(rule['RuleId'])
                request.set_accept_format('json')
                
                self.client.do_action_with_exception(request)
                logger.info(f"成功删除阿里云轻量应用服务器防火墙规则: {rule['RuleId']}")
            except Exception as e:
                logger.error(f"删除轻量应用服务器防火墙规则失败: {e}")
                
    def add_rules(self, instance_id: str, ips: set, port_range: str, description: str):
        """
        添加新的防火墙规则
        
        Args:
            instance_id: 实例ID
            ips: IP地址集合
            port_range: 端口范围
            description: 规则描述(备注)
        """
        for ip in ips:
            try:
                # 将端口范围转换为阿里云轻量应用服务器所需的格式
                # 阿里云轻量应用服务器需要分别添加TCP和UDP规则
                protocols = ['TCP', 'UDP']
                if port_range == "1/65535":
                    port = "-1/-1"  # 全部端口
                else:
                    port = port_range.replace("/", "/")  # 保持原格式
                    
                for protocol in protocols:
                    request = CreateFirewallRuleRequest()
                    request.set_InstanceId(instance_id)
                    request.set_RegionId(self.region_id)
                    request.set_RuleProtocol(protocol)
                    request.set_Port(port)
                    request.set_SourceCidrIp(f"{ip}/32")
                    request.set_Remark(description)  # 使用备注字段存储描述信息
                    request.set_accept_format('json')
                    
                    self.client.do_action_with_exception(request)
                    logger.info(f"成功添加阿里云轻量应用服务器防火墙规则: {ip} ({protocol})")
            except Exception as e:
                logger.error(f"添加轻量应用服务器防火墙规则失败: {e}")
                
    def sync_rules(self, instance_id: str, ips: set, port_range: str = "1/65535", description: str = "Auto-synced rule"):
        """
        同步防火墙规则
        
        Args:
            instance_id: 实例ID
            ips: 当前IP地址集合
            port_range: 端口范围
            description: 规则描述标识
        """
        # 获取现有规则
        existing_rules = self.get_existing_rules(instance_id, description)
        logger.info(f"找到 {len(existing_rules)} 条现有防火墙规则需要删除")
        
        # 删除旧规则
        self.remove_rules(instance_id, existing_rules)
        
        # 添加新规则
        logger.info(f"准备添加 {len(ips)} 条新规则")
        self.add_rules(instance_id, ips, port_range, description)